#!/usr/bin/env python
# -*- coding: UTF-8 -*-
'''
@Project ：V2 
@File    ：hcqs_ga.py
@IDE     ：PyCharm 
@Author  ：郭星
@Date    ：2025/9/22 22:05 
'''
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
HCQS-GA: Heterogeneous-Coverage QC Scheduling solved by GA
----------------------------------------------------------
- 设备异构：每台设备可检测的组分集合不同
- 样品需求：每个样品需要检测的组分集合
- 覆盖决策：样品可由 1..K 台设备共同覆盖（并行）；也可切换为串联模式
- 设备内排序：序列相关 setup（按“方法签名”近似）
- 目标：加权 (Cmax + Tardiness + total Setup)
- GA：二层编码（覆盖/指派 + 设备内顺序），精英保留，单调收敛曲线

Author: you
"""

import random
import numpy as np
from dataclasses import dataclass
from typing import List, Dict, Set, Tuple, Optional

try:
    import matplotlib.pyplot as plt
    HAS_MPL = True
except Exception:
    HAS_MPL = False


# ---------- 问题定义 ----------

@dataclass
class Device:
    id: int
    name: str
    capabilities: Set[str]     # 能检测的组分集合
    base_time: float           # 每个样品基础准备/进出时间
    per_analyte_time: float    # 每个组分的检测时间
    setup_time: float          # 方法签名切换的校准/换型时间

@dataclass
class Sample:
    id: int
    name: str
    requirements: Set[str]     # 需要检测的组分集合
    due: float                 # 交期（目标：总迟期尽量小）

@dataclass
class Task:  # 用于甘特
    device_id: int
    sample_id: int
    start: float
    end: float
    label: str


class Plant:
    def __init__(self, devices: List[Device], samples: List[Sample],
                 max_devices_per_sample: int = 3,
                 chain_mode: str = "parallel"):
        """
        chain_mode:
          - "parallel": 样品可被多台设备并行检测，放行时间 = 所有子检测完成的最大值
          - "serial":   样品在其选择的多台设备上串行，按固定顺序（设备id升序）检测
        """
        self.devices = devices
        self.samples = samples
        self.D = len(devices)
        self.N = len(samples)
        self.kmax = max(1, max_devices_per_sample)
        assert chain_mode in ("parallel", "serial")
        self.chain_mode = chain_mode

    # 设备对某样品的“方法签名”（用于判定是否需要 setup）
    def method_signature(self, sample: Sample, dev: Device) -> Tuple[str, Tuple[str, ...]]:
        # 取样品在该设备上要测的“子组分”集（与cap交集），按字典序形成签名
        sub = tuple(sorted(sample.requirements & dev.capabilities))
        # 也可加入更复杂信息（如灵敏度档位、量程），此处简化
        return ("SIG", sub)

    # 设备对样品的处理时间
    def proc_time(self, sample: Sample, dev: Device) -> float:
        k = len(sample.requirements & dev.capabilities)
        if k == 0:
            return 0.0
        return dev.base_time + k * dev.per_analyte_time


# ---------- 解码（覆盖修复 + 排程 + KPI） ----------

class Decoder:
    def __init__(self, plant: Plant):
        self.P = plant

    def repair_cover(self, cover_bin: np.ndarray) -> np.ndarray:
        """
        修复覆盖：cover_bin[n,d] in {0,1} 表示样品n是否使用设备d。
        若样品未被覆盖其需求，则贪心补设备；若超过kmax，则保留最有用的kmax台。
        返回修复后的 cover_bin。
        """
        N, D = cover_bin.shape
        cov = cover_bin.copy()
        for n, samp in enumerate(self.P.samples):
            need = set(samp.requirements)
            used = [d for d in range(D) if cov[n, d] == 1]
            # 当前覆盖
            covered = set().union(*[self.P.devices[d].capabilities for d in used]) if used else set()
            missing = need - covered
            # 若缺失，则贪心添加能覆盖缺口最多的设备
            if missing:
                cand = list(range(D))
                # 排除已选
                cand = [d for d in cand if d not in used]
                # 贪心添加直到覆盖
                while missing and cand:
                    best = None
                    for d in cand:
                        add = len(self.P.devices[d].capabilities & missing)
                        if add > 0:
                            score = add / (1e-6 + self.P.devices[d].base_time + self.P.devices[d].per_analyte_time * add)
                            if (best is None) or (score > best[0]):
                                best = (score, d, add)
                    if best is None:
                        break
                    _, d, _ = best
                    cov[n, d] = 1
                    used.append(d)
                    covered |= self.P.devices[d].capabilities
                    missing = need - covered
                    cand.remove(d)
            # 若超限，裁剪为最多kmax台，优先保留贡献大的设备
            if np.sum(cov[n]) > self.P.kmax:
                devs = [d for d in range(D) if cov[n, d] == 1]
                devs.sort(key=lambda d: len(self.P.devices[d].capabilities & samp.requirements), reverse=True)
                keep = set(devs[:self.P.kmax])
                for d in devs[self.P.kmax:]:
                    cov[n, d] = 0
            # 再检查：至少覆盖需求
            used = [d for d in range(D) if cov[n, d] == 1]
            covered = set().union(*[self.P.devices[d].capabilities for d in used]) if used else set()
            if not samp.requirements.issubset(covered):
                # 如果仍未覆盖（极端情况），强制选择能覆盖最多缺口的设备直到覆盖
                missing = samp.requirements - covered
                cand = [d for d in range(D) if cov[n, d] == 0]
                while missing and cand:
                    best = max(cand, key=lambda d: len(self.P.devices[d].capabilities & missing))
                    cov[n, best] = 1
                    missing -= self.P.devices[best].capabilities
                    cand.remove(best)
        return cov

    def decode(self, cover_bin: np.ndarray, seq_per_dev: List[List[int]]) -> Tuple[float, Dict[str, float], List[Task]]:
        """
        输入：
          cover_bin[n,d] = 1 表示样品 n 在设备 d 检测
          seq_per_dev[d] = [样品索引...] 为设备 d 的处理顺序（只对 cover=1 的样品有意义）
        输出：
          obj, kpis, tasks（甘特）
        """
        P = self.P
        N, D = cover_bin.shape
        # 设备时间轴
        dev_ready = np.zeros(D, dtype=float)
        last_sig = [None]*D
        tasks: List[Task] = []

        # 若串联模式：为每个样品生成串行顺序（按设备id升序）
        serial_order = {}
        if P.chain_mode == "serial":
            for n in range(N):
                devs = [d for d in range(D) if cover_bin[n, d] == 1]
                devs.sort()
                serial_order[n] = devs

        # 为快速查找每设备的序列中的样品集合
        seq_filtered = [[] for _ in range(D)]
        cover_sets = [set(np.where(cover_bin[:, d] == 1)[0]) for d in range(D)]
        for d in range(D):
            for n in seq_per_dev[d]:
                if n in cover_sets[d]:
                    seq_filtered[d].append(n)

        # 并/串行调度
        # 记录样品在每个设备的结束时间，用于样品放行时刻
        sample_end_times_per_dev = { (n,d): 0.0 for n in range(N) for d in range(D) if cover_bin[n,d]==1 }

        # 串行情况下，需要控制样品在后续设备开始 >= 前一设备结束
        sample_prev_end = { n: 0.0 for n in range(N) }

        # 设备内按给定序列逐一安排
        for d in range(D):
            dev = P.devices[d]
            for n in seq_filtered[d]:
                samp = P.samples[n]
                k = len(samp.requirements & dev.capabilities)
                if k == 0:
                    continue  # 不需要在该设备检测
                pt = P.proc_time(samp, dev)
                # 方法签名与 setup
                sig = P.method_signature(samp, dev)
                setup = 0.0
                if last_sig[d] is None:
                    setup = 0.0
                elif last_sig[d] != sig:
                    setup = dev.setup_time

                # 并行模式：仅受设备就绪约束
                if P.chain_mode == "parallel":
                    st = max(dev_ready[d], 0.0)
                else:
                    # 串联：必须等此样品在该设备的前一环节结束
                    # 如果该样品只在该设备检测（单一），等价于并行
                    prev_needed = serial_order[n]
                    idx = prev_needed.index(d)
                    if idx == 0:
                        st = max(dev_ready[d], 0.0)
                    else:
                        st = max(dev_ready[d], sample_prev_end[n])

                st += setup
                en = st + pt

                tasks.append(Task(d, n, st, en, f"{dev.name}"))
                dev_ready[d] = en
                last_sig[d] = sig
                sample_end_times_per_dev[(n,d)] = en
                if P.chain_mode == "serial":
                    sample_prev_end[n] = en

        # 样品放行时刻
        release = np.zeros(N, dtype=float)
        for n in range(N):
            used = [d for d in range(D) if cover_bin[n, d] == 1]
            if not used:
                release[n] = 0.0
            else:
                release[n] = max(sample_end_times_per_dev[(n,d)] for d in used)

        # KPI
        Cmax = float(np.max(release)) if len(release)>0 else 0.0
        tard = 0.0
        for n, samp in enumerate(P.samples):
            tard += max(0.0, release[n] - samp.due)
        # 统计总 setup
        total_setup = 0.0
        for d in range(D):
            dev = P.devices[d]
            last = None
            for n in seq_filtered[d]:
                sig = P.method_signature(P.samples[n], dev)
                if last is not None and sig != last:
                    total_setup += dev.setup_time
                last = sig

        kpis = dict(Cmax=Cmax, Tardiness=tard, Setup=total_setup)
        # 目标加权（可调整权重）
        obj = 1.0*Cmax + 2.0*tard + 0.2*total_setup
        return obj, kpis, tasks


# ---------- 遗传算法（精英保留 + 单调收敛） ----------

class GA:
    """
    染色体：
      - 覆盖层：cover_bin (N x D) 0/1（限制每样品 <= kmax）
      - 序列层：每个设备一个样品序列（全排列）；解码时忽略未覆盖的样品
    进化：
      - 选择：锦标赛
      - 交叉：覆盖层单点/均匀混合；序列层次序交叉（OX）
      - 变异：覆盖位翻转（并修复）；序列交换/移位
      - 精英保留：top E 直接复制
    """
    def __init__(self, plant: Plant, pop=60, gens=300, pc=0.9, pm=0.3, elite=None, seed=0):
        self.P = plant
        self.N = len(plant.samples)
        self.D = len(plant.devices)
        self.pop = pop
        self.gens = gens
        self.pc = pc
        self.pm = pm
        self.elite = elite if elite is not None else max(2, pop//10)
        random.seed(seed); np.random.seed(seed)
        self.dec = Decoder(plant)

    def rand_ind(self):
        # 覆盖：每个样品随机选择 1..kmax 台能覆盖的设备（若随机不到覆盖，后续 repair）
        cover = np.zeros((self.N, self.D), dtype=np.int8)
        for n, samp in enumerate(self.P.samples):
            # 候选设备（能覆盖其任一组分即可）
            cand = [d for d in range(self.D) if len(self.P.devices[d].capabilities & samp.requirements) > 0]
            k = np.random.randint(1, min(self.P.kmax, max(1, len(cand)))+1)
            pick = list(np.random.choice(cand, size=k, replace=False))
            cover[n, pick] = 1
        cover = self.dec.repair_cover(cover)
        # 序列：每个设备随机一个样品排列
        seq = [list(np.random.permutation(self.N)) for _ in range(self.D)]
        return (cover, seq)

    # 序列的 OX 交叉
    def ox(self, a: List[int], b: List[int]) -> List[int]:
        n=len(a); i,j=sorted(np.random.choice(n,2,replace=False))
        child=[None]*n; child[i:j]=a[i:j]
        pos=j
        for x in b:
            if x not in child:
                if pos>=n: pos=0
                child[pos]=x; pos+=1
        return child

    def crossover(self, p1, p2):
        c1_cov = p1[0].copy(); c2_cov = p2[0].copy()
        if np.random.rand()<0.5:
            # 覆盖层均匀交叉
            mask = (np.random.rand(self.N, self.D) < 0.5)
            c1_cov[mask] = p2[0][mask]
            c2_cov[mask] = p1[0][mask]
        else:
            # 单点按样品行
            cut = np.random.randint(1, self.N)
            c1_cov[:cut,:] = p1[0][:cut,:]; c1_cov[cut:,:] = p2[0][cut:,:]
            c2_cov[:cut,:] = p2[0][:cut,:]; c2_cov[cut:,:] = p1[0][cut:,:]
        # 序列层 OX
        c1_seq=[]; c2_seq=[]
        for d in range(self.D):
            c1_seq.append(self.ox(p1[1][d], p2[1][d]))
            c2_seq.append(self.ox(p2[1][d], p1[1][d]))
        return (c1_cov, c1_seq), (c2_cov, c2_seq)

    def mutate(self, ind, pm_cov=0.2, pm_seq=0.2):
        cov, seq = ind
        # 覆盖位翻转
        if np.random.rand()<pm_cov:
            n = np.random.randint(self.N); d = np.random.randint(self.D)
            cov[n,d] = 1 - cov[n,d]
        # 行内裁剪（不超过 kmax）
        for n in range(self.N):
            if cov[n].sum() > self.P.kmax:
                ones = np.where(cov[n]==1)[0].tolist()
                while len(ones) > self.P.kmax:
                    d = random.choice(ones); cov[n,d]=0; ones.remove(d)
        # 序列扰动：交换或移位
        if np.random.rand()<pm_seq:
            d = np.random.randint(self.D)
            if len(seq[d])>=2:
                if np.random.rand()<0.5:
                    i,j = sorted(np.random.choice(len(seq[d]),2,replace=False))
                    seq[d][i], seq[d][j] = seq[d][j], seq[d][i]
                else:
                    i = np.random.randint(len(seq[d]))
                    val = seq[d].pop(i)
                    j = np.random.randint(len(seq[d])+1)
                    seq[d].insert(j, val)
        return (cov, seq)

    def fitness(self, ind):
        cov, seq = ind
        cov = self.dec.repair_cover(cov)
        obj, kpis, _ = self.dec.decode(cov, seq)
        return obj, kpis

    def tournament(self, pop_fit, k=2):
        idx = np.random.choice(len(pop_fit), k, replace=False)
        best = min(idx, key=lambda i: pop_fit[i][1])  # 按 obj
        return pop_fit[best][0]

    def run(self, verbose=True):
        # 初始化
        pop = [self.rand_ind() for _ in range(self.pop)]
        fits = []
        for ind in pop:
            obj, kpis = self.fitness(ind)
            fits.append((ind, obj, kpis))
        best_so_far = min(fits, key=lambda x: x[1])
        best_hist = [best_so_far[1]]

        for g in range(1, self.gens+1):
            # 精英
            fits.sort(key=lambda x: x[1])
            elites = fits[:self.elite]
            new_pop = [e[0] for e in elites]

            # 产生子代
            while len(new_pop) < self.pop:
                p1 = self.tournament(fits, k=2)
                p2 = self.tournament(fits, k=2)
                c1, c2 = p1, p2
                if np.random.rand()<self.pc:
                    c1, c2 = self.crossover(p1, p2)
                if np.random.rand()<self.pm:
                    c1 = self.mutate(c1)
                if np.random.rand()<self.pm:
                    c2 = self.mutate(c2)
                new_pop.extend([c1, c2])
            pop = new_pop[:self.pop]

            # 评估
            fits = []
            for ind in pop:
                obj, kpis = self.fitness(ind)
                fits.append((ind, obj, kpis))
            cur_best = min(fits, key=lambda x: x[1])
            if cur_best[1] < best_so_far[1]:
                best_so_far = cur_best
            best_hist.append(best_so_far[1])
            if verbose and g % 20 == 0:
                print(f"[GA] Gen {g:4d}  Best obj so far = {best_so_far[1]:.3f}")

        # 最优解细解码，拿任务甘特
        best_cov, best_seq = best_so_far[0]
        best_cov = self.dec.repair_cover(best_cov)
        _, best_kpis, tasks = self.dec.decode(best_cov, best_seq)
        return (best_cov, best_seq), best_so_far[1], best_kpis, best_hist, tasks


# ---------- 可视化 ----------

def plot_convergence(hist: List[float], title="Convergence (Monotone)"):
    if not HAS_MPL:
        print("(matplotlib not installed; skip convergence plot)"); return
    plt.figure(figsize=(6,4))
    plt.plot(hist, lw=2)
    plt.xlabel("Generation"); plt.ylabel("Best-so-far Objective")
    plt.title(title); plt.grid(True, alpha=0.3); plt.tight_layout(); plt.show()

def plot_gantt(tasks: List[Task], plant: Plant, title="QC Gantt"):
    if not HAS_MPL or not tasks:
        if not HAS_MPL: print("(matplotlib not installed; skip gantt)")
        return
    D = len(plant.devices)
    fig, ax = plt.subplots(figsize=(12, 0.8*D + 2))
    by_dev: Dict[int, List[Task]] = {}
    for t in tasks:
        by_dev.setdefault(t.device_id, []).append(t)
    for d in range(D):
        row = d
        dev = plant.devices[d]
        lst = sorted(by_dev.get(d, []), key=lambda x:x.start)
        for t in lst:
            ax.barh(row, t.end - t.start, left=t.start, height=0.8)
            ax.text((t.start+t.end)/2, row, f"S{t.sample_id}", ha='center', va='center', fontsize=8, color='white')
    ax.set_yticks(range(D)); ax.set_yticklabels([dev.name for dev in plant.devices])
    ax.set_xlabel("Time"); ax.set_title(title)
    ax.grid(axis='x', alpha=0.3)
    plt.tight_layout(); plt.show()


# ---------- 数据生成（示例） ----------

def build_demo(seed=0, mode="parallel"):
    random.seed(seed); np.random.seed(seed)

    # 设备：三台，覆盖不同组分
    devs = [
        Device(0, "E1", {"A","B","C"}, base_time=0.4, per_analyte_time=0.6, setup_time=0.5),
        Device(1, "E2", {"D","E","F"}, base_time=0.5, per_analyte_time=0.7, setup_time=0.6),
        Device(2, "E3", {"G","H","I"}, base_time=0.4, per_analyte_time=0.5, setup_time=0.4),
        # 你可以再加多台设备，能力可以重叠
    ]

    # 样品：生成一批需要多组分的产品（示意）
    all_g = ["A","B","C","D","E","F","G","H","I"]
    samples=[]
    N=18
    for i in range(N):
        # 随机抽 2~4 个组分需求（更接近真实混合）
        req = set(np.random.choice(all_g, size=np.random.randint(2,5), replace=False))
        # 交期分布
        due = np.random.uniform(6.0, 20.0)
        samples.append(Sample(i, f"P{i}", req, due))

    plant = Plant(devs, samples, max_devices_per_sample=3, chain_mode=mode)
    return plant


# ---------- 主程序 ----------

def main():
    # 切换并/串模式：parallel 或 serial
    plant = build_demo(seed=3, mode="parallel")

    ga = GA(plant, pop=80, gens=240, pc=0.9, pm=0.35, elite=10, seed=1)
    (best_cov, best_seq), best_obj, kpis, hist, tasks = ga.run(verbose=True)

    print("\n== KPIs ==")
    for k,v in kpis.items():
        print(f"{k:>10s}: {v:.3f}")
    print(f"{'Objective':>10s}: {best_obj:.3f}")

    # 打印部分覆盖决策（前6个样品）
    print("\nCoverage (sample -> devices):")
    for n in range(min(6, plant.N)):
        devs = [plant.devices[d].name for d in range(plant.D) if best_cov[n,d]==1]
        print(f"S{n:<2d} -> {devs}")

    plot_convergence(hist, title=f"HCQS-GA ({plant.chain_mode})")
    plot_gantt(tasks, plant, title=f"QC Gantt ({plant.chain_mode})")

if __name__ == "__main__":
    main()
